from state import EditorState
from langchain_core.prompts import ChatPromptTemplate
from model.zhipu import get_model
from utils.get_json import get_json

prompt = ChatPromptTemplate.from_template("""
你是文章总结编撰,你需要利用你的专业知识为提交对文档进行内容整合.
 - 期撰写的主题为{title}
 - 文章中的所有信息为{paragraphs}
 - 你需要返回一个完整的文本
""")


def generate_summary(state: EditorState):
    title = state['title']
    paragraphs = state['paragraphs']

    model = get_model()
    chains = prompt | model
    response = chains.invoke({"title": title, "paragraphs": paragraphs})

    print("============ Summary Agent ============")
    print(response.content)
    # comment_info = get_json(response.content)

    return {
        "who": "comment",
        "title": title,  # 制定标题
        "search": "",  # 搜集信息
        "outlines": {},  # 大纲搜集信息 三级标题
        "paragraphs": paragraphs,  # 已完成的段落信息
        "write_submit": "",
        "check_idea": {},  # 校对意见
        "comment_idea": {}  # 评论意见
    }


def commit_to_write_agent(state: EditorState):
    check_idea = state['check_idea']
    if check_idea['pass']:
        return "comment_agent"
    else:
        return "write_agent"

# if __name__ == '__main__':
#     inputs = {
#         "who": "",
#         "title": "中国的近代历史",  # 制定标题
#         "search": "",  # 搜集信息
#         "outlines": {},  # 大纲搜集信息 三级标题
#         "paragraphs": [],  # 段落信息
#         "check_idea": {},  # 校对意见
#         "comment_idea": {}  # 评论意见
#     }
#
#     outputs = generate_title(inputs)
#     print(outputs)
#     print(outputs['title'])
